import pandas as pd
import logging
import json
import os
from typing import Dict, List, Optional, Tuple, Any

from .casualEdge import CausalityEdge
from .casualGraph import CausalGraph
from source.data_processor.utils.experiment_data import ExperimentData

# 配置日志
logger = logging.getLogger(__name__)

class CompositeCausalGraph:
    """
    综合因果图，用于管理故障前和故障后的因果图，并提供综合分析功能
    """
    
    def __init__(self, uri, user, password, cache_dir=None, batch_size=2000, min_data_points=1, output_dir: str = None):
        """
        初始化综合因果图
        
        参数:
            uri: Neo4j数据库URI
            user: Neo4j用户名
            password: Neo4j密码
            cache_dir: 缓存目录，用于存储PKL文件
            batch_size: 批处理大小
            min_data_points: 最小数据点数，用于因果分析
            output_dir: 输出目录，用于存储各阶段图的JSON状态文件
        """
        # 基础配置
        self.uri = uri
        self.user = user
        self.password = password
        self.cache_dir = cache_dir
        self.output_dir = output_dir
        self.batch_size = batch_size
        self.min_data_points = min_data_points
        
        # 存储故障前后的图
        self.pre_anomaly_graph = None
        self.post_anomaly_graph = None
        
        # 存储综合的因果边
        self.combined_edges = []
        
        # ART特征提取结果
        self.art_features = None
        
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        # 关闭子图的数据库连接
        self._close_graphs()
            
    def _close_graphs(self):
        """关闭所有子图的数据库连接"""
        if self.pre_anomaly_graph:
            self.pre_anomaly_graph.close()
        if self.post_anomaly_graph:
            self.post_anomaly_graph.close()
    
    def build_composite_graph(self, experiment_data: ExperimentData, metric_result=None, trace_result=None, extract_art_features=True):
        """
        构建故障前后的综合因果图
        
        参数:
            experiment_data: 实验数据对象
            metric_result: 指标异常检测结果，用于筛选异常指标
            trace_result: 链路异常检测结果，用于筛选相关的Interface和Container
            extract_art_features: 是否提取ART特征
            
        返回:
            bool: 构建成功返回True，否则返回False
        """
        try:
            # 1. 获取基本信息
            graph_id = experiment_data.anomaly_id
            logger.info(f"开始构建综合因果图, graph_id={graph_id}")
            
            # 2. 获取正常和异常时间段
            phases = self._prepare_time_phases(experiment_data)
            if not phases:
                return False
            
            # 3. 构建正常和异常阶段的因果图
            pre_success = self._build_phase_graph(
                "pre_anomaly", 
                experiment_data, 
                phases["pre_anomaly"], 
                metric_result, 
                trace_result
            )
            
            # post_success = self._build_phase_graph(
            #     "post_anomaly", 
            #     experiment_data, 
            #     phases["post_anomaly"], 
            #     metric_result, 
            #     trace_result
            # )
            
            # 4. 合并因果边
            self._combine_causal_edges()
            
            # 5. 汇报结果
            self._report_results(graph_id)
            
            return pre_success or post_success
            
        except Exception as e:
            logger.error(f"构建综合因果图时发生错误: {e}")
            import traceback
            logger.error(traceback.format_exc())
            return False
    
    def _prepare_time_phases(self, experiment_data: ExperimentData) -> Optional[Dict[str, Tuple[Any, Any]]]:
        """
        准备时间阶段数据
        
        参数:
            experiment_data: 实验数据对象
            
        返回:
            Dict: 包含各阶段时间范围的字典，如果获取失败则返回None
        """
        try:
            normal_time_range = experiment_data.get_normal_time_range()
            anomaly_time_range = experiment_data.get_anomaly_time_range()
            
            logger.info(f"获取到正常时间段: {normal_time_range[0]} 至 {normal_time_range[1]}")
            logger.info(f"获取到异常时间段: {anomaly_time_range[0]} 至 {anomaly_time_range[1]}")
            
            return {
                "pre_anomaly": normal_time_range,
                "post_anomaly": anomaly_time_range
            }
        except ValueError as e:
            logger.error(f"获取时间段失败: {e}")
            return None
    
    def _create_causal_graph(self, phase: str, graph_id: str) -> CausalGraph:
        """
        创建特定阶段的因果图实例
        
        参数:
            phase: 阶段标识，如 'pre_anomaly' 或 'post_anomaly'
            graph_id: 图ID
            
        返回:
            CausalGraph: 创建的因果图实例
        """
        phase_graph_id = f"{graph_id}_{phase}"
        
        return CausalGraph(
            uri=self.uri,
            user=self.user,
            password=self.password,
            cache_dir=self.cache_dir,
            output_dir=self.output_dir,
            batch_size=self.batch_size,
            min_data_points=self.min_data_points,
            phase=phase,
            active_graph_id=phase_graph_id
        )
    
    def _build_phase_graph(self, phase: str, experiment_data: ExperimentData, time_range: Tuple[Any, Any], metric_result=None, trace_result=None) -> bool:
        """
        构建特定阶段的因果图
        
        参数:
            phase: 阶段标识，如 'pre_anomaly' 或 'post_anomaly'
            experiment_data: 实验数据对象
            time_range: 时间范围元组
            metric_result: 指标异常检测结果
            trace_result: 链路异常检测结果
            
        返回:
            bool: 构建成功返回True，否则返回False
        """
        # 创建阶段图
        graph = self._create_causal_graph(phase, experiment_data.anomaly_id)
        
        # 存储图实例
        if phase == "pre_anomaly":
            self.pre_anomaly_graph = graph
        else:
            self.post_anomaly_graph = graph
        
        # 构建因果图
        success = graph.build_causal_graph(
            experiment_data,
            metric_result,
            trace_result=trace_result,
            time_range=time_range
        )
        
        return success
    
    def _combine_causal_edges(self):
        """合并两个阶段的因果边"""
        pre_edges = self.pre_anomaly_graph.causal_edges if self.pre_anomaly_graph else []
        post_edges = self.post_anomaly_graph.causal_edges if self.post_anomaly_graph else []
        
        self.combined_edges = pre_edges + post_edges
    
    def _report_results(self, graph_id: str):
        """报告构建结果"""
        pre_count = len(self.pre_anomaly_graph.causal_edges) if self.pre_anomaly_graph else 0
        post_count = len(self.post_anomaly_graph.causal_edges) if self.post_anomaly_graph else 0
        total_count = len(self.combined_edges)
        
        logger.info(f"综合因果图构建完成, graph_id={graph_id}")
        logger.info(f"因果关系数量: 总计 {total_count} (故障前: {pre_count}, 故障后: {post_count})")
    
    def get_causal_analysis_results(self) -> Dict[str, List[CausalityEdge]]:
        """
        获取因果分析结果
        
        返回:
            dict: 包含 'all', 'pre_anomaly', 'post_anomaly' 三个键的字典，值为对应的因果边列表
        """
        pre_edges = self.pre_anomaly_graph.causal_edges if self.pre_anomaly_graph else []
        post_edges = self.post_anomaly_graph.causal_edges if self.post_anomaly_graph else []
        
        return {
            'all': self.combined_edges,
            'pre_anomaly': pre_edges,
            'post_anomaly': post_edges
        }
